import rclpy
from rclpy.node import Node
# import pupil_apriltags as apriltag
import apriltag
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
import cv2
import math
from sls_msgs.msg import SLSApirltag

class ApirltagDetect(Node):
    def __init__(self,source):
        super().__init__('apirltag_detect_node')
        self.sub = self.create_subscription(Image,'img_raw',self.ApirltagDetectCallback,10)
        self.pub_det = self.create_publisher(Image,'img_apirltag_detect',10)
        self.timer_det = self.create_timer(0.03,self.timerDetcallback)

        self.pub_status = self.create_publisher(SLSApirltag,'img_apirltag_status',10)
        self.timer_status = self.create_timer(0.03,self.timerStatuscallback)
        self.cap = cv2.VideoCapture(source)
        if not self.cap.isOpened():
            self.get_logger().error('SLS: Failed to open camera, please ensure your camera source is 2!')

        self.cv_bridge = CvBridge()
        # self.detector = apriltag.Detector(families = 'tag36h11')
        self.detector = apriltag.Detector()

    def ApirltagDetectCallback(self,image):
        self.frame = self.cv_bridge.imgmsg_to_cv2(image)
        # self.gray_frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY)

    def timerDetcallback(self):
        try:
            self.detect(screem=None)
            self.pub_det.publish(self.cv_bridge.cv2_to_imgmsg(self.frame,'mono8'))
        except:
            self.get_logger().warning('no image data receive')
            pass
        
    def timerStatuscallback(self):
        try:
            coordinate = SLSApirltag()
            coordinate.x = self.cX
            coordinate.y = self.cY
            coordinate.angle = self.angle
            coordinate.isdetect = self.isdetect
            self.pub_status.publish(coordinate)
        except:
            self.get_logger().warning('no coordinate')
            pass
    def detect(self, msgs: bool = True , screem: bool = True)->tuple[int, int , float]:
        '''
        Args:
            msgs (bool): If true, print(c,y,z) in terminal
            screem (bool): If true, print(c,y,z) in screem
        return:
            tuple[int, int , int] | None: If detected, return (cX,cY,angle), else return None
        '''
        try:
            results = self.detector.detect(self.frame)
            if results != []:
                for result in results:
                    cv2.line(self.frame,tuple(result.corners[0].astype(int)),
                                tuple(result.corners[1].astype(int)),(0, 255, 0),10)
                    cv2.line(self.frame,tuple(result.corners[1].astype(int)),
                                tuple(result.corners[2].astype(int)),(0, 255, 0),10)
                    cv2.line(self.frame,tuple(result.corners[2].astype(int)),
                                tuple(result.corners[3].astype(int)),(0, 255, 0),10)
                    cv2.line(self.frame,tuple(result.corners[3].astype(int)),
                                tuple(result.corners[0].astype(int)),(0, 255, 0),10)
                    (self.cX, self.cY) = (int(result.center[0]), int(result.center[1]))
                    cv2.circle(self.frame, (self.cX, self.cY), 10, (0, 0, 255), -1)

                    self.angle = math.atan((result.corners[0][1]-result.corners[3][1])/
                                        (result.corners[0][0]-result.corners[3][0]))*57.3
                    # print(result.corners[0][1],result.corners[3][1])
                    # if(result.corners[0][1]<result.corners[3][1]):
                    #     angle += 180
                    self.isdetect =True
                    if screem == True:
                        cv2.imshow("video", self.frame)
                    if msgs == True:
                        print(f'center = {(self.cX, self.cY)}')
                        print(f'angle = {self.angle}')
                    return self.cX,self.cY,self.angle
            else:
                self.isdetect =False
                if msgs == True:
                    self.get_logger().warning('no detect')
                if screem == True:
                    cv2.imshow("video", self.frame)
                return None
        except:
            self.isdetect =False

def main(args=None):
    rclpy.init(args=args)
    node = ApirltagDetect()
    rclpy.spin(node)
    node.destroy_node(node)
    rclpy.shutdown()